import os
import shutil
import subprocess


def modify_upk_file(file_path, output_dir):
    # 读取文件内容
    with open(file_path, 'rb') as f:
        data = bytearray(f.read())

    # 修改前4字节
    data[0:4] = b'\xFF\xA8\x27\x3E'

    # 修改第5字节
    data[4] = (data[4] + 8) % 256

    # 生成输出文件路径
    file_name = os.path.basename(file_path)
    modified_file_path = os.path.join(output_dir, file_name.replace('.upk', '_modified.upk'))

    # 写回修改后的文件
    with open(modified_file_path, 'wb') as f:
        f.write(data)

    return modified_file_path


def run_external_programs(modified_file_path):
    print(f"Processing {modified_file_path} with UnPack.exe")
    subprocess.run(['UnPack.exe', modified_file_path])

    # 创建 tempUpk 目录
    temp_dir = 'tempUpk'
    os.makedirs(temp_dir, exist_ok=True)

    # 查找并移动 _ucp 文件，在当前脚本所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    for file in os.listdir(current_dir):
        if file.endswith('_ucp.upk'):
            ucp_file_path = os.path.join(current_dir, file)
            destination_path = os.path.join(temp_dir, file)
            shutil.move(ucp_file_path, destination_path)
            absolute_ucp_path = os.path.abspath(destination_path)

            print(f"Repairing {absolute_ucp_path} with Repair.exe")
            process = subprocess.Popen(['Repair.exe', modified_file_path], stdin=subprocess.PIPE)
            process.communicate(input=f'0\n{absolute_ucp_path}\n125\n'.encode())
            process.kill()

            # 修改 ucp 文件内容
            with open(modified_file_path, 'r+b') as f:
                f.seek(0x04)
                f.write(b'\x2C\x03')

            print(f"Decompressing {modified_file_path} with decompress.exe")
            subprocess.run(['decompress.exe', modified_file_path])


def process_upk_files(input_dir, output_dir):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 遍历输入目录中的 .upk 文件
    for root, _, files in os.walk(input_dir):
        for file in files:
            if file.endswith('.upk'):
                file_path = os.path.join(root, file)
                print(f"Modifying {file_path}")
                modified_file_path = modify_upk_file(file_path, output_dir)
                run_external_programs(modified_file_path)
                print(f"Completed processing {file_path}\n")


# 使用方法
input_directory = 'upkFiles'
output_directory = 'outUpk'
process_upk_files(input_directory, output_directory)
